#ifndef _REMOVE_BUBBLES_CORRECTOR_H_
#define _REMOVE_BUBBLES_CORRECTOR_H_

#include "SequenceGraph.h"
#include "FibHeap.h"
#include <set>
#include "GraphErrorCorrector.h"

using std::set;

class SimpleMergeCorrector: public AbstractGraphCorrector {
public:
	SimpleMergeCorrector(SequenceGraph* seqGraph): AbstractGraphCorrector(seqGraph) {}
	virtual SequenceGraph* correctGraph();
};

class DissectBubblesCorrector: public AbstractGraphCorrector {	
private:
	class TraceNode: public AbstractHeapNode {
	public:
		TraceNode( GID_T seqVID, TraceNode* p = NULL): AbstractHeapNode(), parent(p), seqVertexID(seqVID), dist(0) {}
		// default copy constructor is ok.
		TraceNode* parent;
		GID_T seqVertexID;
		vector<TraceNode> children;
		int dist;
		inline static int cmp_func(void* node1, void* node2)
		{
			assert(node1 != NULL && node2 != NULL);
			TraceNode* ele1 = (TraceNode*) node1;
			TraceNode* ele2 = (TraceNode*) node2;
			if (ele1 == ele2 || ele1->dist == ele2->dist) {
				return 0;
			} else if (ele1->dist > ele2->dist) {
				return 1;
			} else {
				return -1;
			}
		}
	};

	class TraceRecord {
	private:
		typedef map<GID_T, vector<vector<TraceNode*> > > map_type;
		typedef map_type::iterator iterator_type;
		map_type m_map;
		int treeCount;
	public:
		TraceRecord(int treeCnt): treeCount(treeCnt) {}
		/**
		 * Record sequence vertex is traced in the node of which tree.
		 * @param seqVID the ID of sequence vertex
		 * @param treeID the ID of trace tree
		 * @param node the node on the tree
		 */
		void record(GID_T seqVID, int treeID, TraceNode* node);

		/**
		 * Get possible paths that are recorded in other trees.
		 * If there are trees that are terminated, its record will be removed
		 * and thus not be returned.
		 *
		 * @param node trace node containing id and dist of the vertex to query
		 * @param treeID id of tree not meant to include
		 * @param absMismatch maximal mismatch allowed
		 * @param pruneBound records whose dist of less than this value will never use
		 *
		 * @return the shared end of other possible paths with tree ID
		 */
		vector<pair<int, TraceNode*> > otherPossiblePaths(const TraceNode& node, int treeID, int absMismatch, int pruneBound);

		void removeRecord(int treeID);

		~TraceRecord();
	};

	class DistNodeMap {
	private:
		typedef map<int, vector<vector<DissectBubblesCorrector::TraceNode*> > > map_type;
		typedef map_type::iterator iterator_type;
		map_type m_map;
		int m_treeCnt;
	public:
		DistNodeMap(int treeCnt): m_treeCnt(treeCnt) {}
		/**
		 * Gets nearby nodes in range [dist - absMismatch, dist + absMistmatch]
		 * and remove anything in range [0, dist - absMismatch)
		 * @param dist 
		 * @param absMismatch
		 * @param treeID the id of tree to exclude
		 * @return 
		 */

		~DistNodeMap();

		vector<TraceNode*> vicinities(int dist, int absMismatch, int treeID);

		void insert(int treeID, TraceNode* node);
	};

	class CandidateNode: public AbstractHeapNode {
	public:
		CandidateNode( GID_T seqVID, int branchID ): AbstractHeapNode(), seqVertexID(seqVID), branch(branchID) , dist(0) {}
		// default copy constructor is ok.
		GID_T seqVertexID;
		int branch;
		int dist;
		vector<CandidateNode*> parents;
		inline static int cmp_func(void* node1, void* node2)
		{
			assert(node1 != NULL && node2 != NULL);
			CandidateNode* ele1 = (CandidateNode*) node1;
			CandidateNode* ele2 = (CandidateNode*) node2;
			if (ele1 == ele2 || ele1->dist == ele2->dist) {
				return 0;
			} else if (ele1->dist > ele2->dist) {
				return 1;
			} else {
				return -1;
			}
		}
	};

	int m_searchBound;
	int m_mismatchUpperBound; //< max mismatch ever allowed
	int m_mismatchLowerBound; //< mismatch less the this is always allowed
	double m_relativeMismatch;
	unsigned char* m_visitRecord;
	unsigned char* m_traceRecord;
	
	// for decompressed kmers
	BitDeque p1content;
	BitDeque p2content;
	int* distMatrix; //< used to calculate edit distance

	SimpleMergeCorrector m_simpleMerger;

	inline bool isVisited(GID_T id) { id = abs(id); return (m_visitRecord[id >> 3] >> (id % 8)) & 1; }
	inline void setVisited(GID_T id) { id = abs(id); m_visitRecord[id >> 3] |= (1 << (id % 8)); }
	inline bool isTraced(GID_T id) { id = abs(id); return (m_traceRecord[id >> 3] >> (id % 8)) & 1; }
	inline void setTraced(GID_T id) { id = abs(id); m_traceRecord[id >> 3] |= (1 << (id % 8)); }
	//bool isTerminated(const vector<FibHeap>& queues, int tid);
	SequenceVertex* nextRoot(); //find a root to bfs the graph, positive one has no predecessor, twin one if the positive one has no successor.
	void dissectBubble(SequenceVertex* root);
	void tryResolve(SequenceVertex* end);
	void tryResolve(const vector<SequenceEdge>& allpreds, int one, int another, SequenceVertex* end, vector<bool>& detached, const set<GID_T>& candidates);
	bool contain(const SequenceEdge& ele, const vector<SequenceEdge>& collection) const;
	/**
	 *
	 * @param queues the heaps keeping shortest branch on the top of the trees
	 *
	 * @return id of lowest tree or -1 if all trees cannot grow (heap empty)
	 */
	int chooseLowestTree( const vector<FibHeap>& queues );

	/**
	 * Extract bases of the path. The first k-1 and last k-1 bases are truncated
	 * for they must be the same on both paths.
	 * @param pathVertices 
	 * @param [out] pathContent 
	 */
	void recoverPath(const vector<SequenceVertex*>& pathVertices, BitDeque& pathContent);
	
	/**
	 * TODO: an SSE2 vectorization may help to improve speed; space can further be reduced to
	 * 2(2*mismatch)
	 * @param str1 
	 * @param str2 
	 * @return 
	 */
	bool areSimilar(const BitDeque& str1, const BitDeque& str2, int threshold);
	inline int getArrIdx(char ch);

	int generateCandidates(SequenceVertex* pred1, SequenceVertex* pred2, SequenceVertex* end, set<GID_T>& candidates);
	int chooseLowestBranch( const vector<FibHeap>& queues );

	bool shouldGrow( TraceNode* curNode, int curTreeId, DistNodeMap& distMap );
	void recoverPathCompleteHead(const vector<SequenceVertex*>& pathVertices, BitDeque& pathContent);
	void cleanTrees( vector<TraceNode>& roots );
public:
	static int replace_penalty;
	static int indel_penalty;
	/**
	 * Constructor.
	 * If two paths forms a bubble and are similar enough, they will be merged.
	 *
	 * @param seqGraph the sequence graph
	 * @param searchBound distance limitation for searching bubbles
	 * @param mismatchUpperBound bigest mismatch allowed to merge two paths
	 * @param relativeMismatch bigest mismatch rate allowed to merge two paths
	 *
	 * @return 
	 */
	DissectBubblesCorrector(SequenceGraph* seqGraph, int searchBound, int mismatchUpperBound, int mismatchLowerBound, double relativeMismatch);
	~DissectBubblesCorrector();
	SequenceGraph* correctGraph();
};

#endif